# -*- coding: UTF-8 -*-
import os
import threading
import yaml
import numpy as np
import rospy
import cv2
from sensor_msgs.msg import Image
from sensor_msgs.msg import CompressedImage
from sensor_msgs.msg import Imu
from cv_bridge import CvBridge
import module.IMU_Optimize
import module.lowlight_enhancement

#pip install opencv-python==3.4.3.18


bridge = CvBridge()
i=0
imuList=[]
imageList=[]
imuCheckingList=[]
imgProcessingList=[]

imuSendList=[]
imageSendList=[]

imu_buf=threading.Lock()
img_buf=threading.Lock()
processing=threading.Lock()
processingImage=threading.Lock()

data_train=None
imu_topic=None
imu_pub=None
image_topic=None
image_pub=None

def initParameter():
    global  imu_topic,image_topic,imu_pub,image_pub
    fs = open(os.path.join("../config", "parameter.yaml"), encoding="UTF-8")
    datas = yaml.load(fs, Loader=yaml.FullLoader)
    imu_topic=datas['imu_topic']
    image_topic=datas["image_topic"]
    imu_pub=datas["imu_pub"]
    image_pub=datas["image_pub"]


def get_imu_record(path):
    list_files = []
    if os.path.exists(path):
        # 遍历文件夹，找到.npy文件
        for root, ds, fs in os.walk(path):
            for f in fs:
                if f.endswith('.npy'):
                    fullname = os.path.join(root, f)
                    list_files.append(fullname)
        for file in list_files:
            lab = os.path.split(os.path.split(file)[0])[-1]
    else:
        print('this path not exist')
    return list_files



def handleIMU(data):
    imu_buf.acquire()
    imuList.append(data)
    imu_buf.release()

def handleImage(data):
    img_buf.acquire()
    imageList.append(data)
    img_buf.release()

def handleCompressedImage(data):
    img_buf.acquire()
    imageList.append(data)
    img_buf.release()

def initialPublishImuByGroup():
    global imuList,imuSendList
    imu_buf.acquire()
    imuSendList=imuList
    imuList=[]
    imu_buf.release()
    #imuSendList.reverse()
    while len(imuSendList)>0:
        imupub.publish(imuSendList.pop())
        imurate.sleep()


def initialPublishImgByGroup():
    global imageSendList,imgProcessingList
    img_buf.acquire()
    imageSendList=imgProcessingList
    imgProcessingList=[]
    img_buf.release()
    imageSendList.reverse()
    while len(imageSendList)>0:
        imgpub.publish(imageSendList.pop())
        imgrate.sleep()

def imgProcess():
    global imageList,imgProcessingList
    img_buf.acquire()
    processingImage.acquire()
    imgProcessingList=imageList
    imageList=[]
    img_buf.release()
    template=[]
    for data in imgProcessingList:
        cv_image = bridge.imgmsg_to_cv2(data, desired_encoding='bgr8')
        from PIL import Image
        frame = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
        result=module.lowlight_enhancement.lowlight(frame)
        img = cv2.cvtColor(np.asarray(result), cv2.COLOR_RGB2BGR)
        image_message = bridge.cv2_to_imgmsg(img, encoding="bgr8")
        template.append(image_message)
    for frame,target in zip(imgProcessingList,template):
        frame.data=target.data
        frame.encoding=target.encoding
        frame.height=target.height
        frame.width=target.width
        frame.step=target.step
    processingImage.release()

def imgCompressedProcess():
    global imageList,imgProcessingList
    img_buf.acquire()
    processingImage.acquire()
    imgProcessingList=imageList
    imageList=[]
    img_buf.release()
    template=[]
    for im in imgProcessingList:
        np_arr = np.fromstring(im.data, np.uint8)
        cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        from PIL import Image as Im
        frame = Im.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
        result=module.lowlight_enhancement.lowlight(frame)
        img = cv2.cvtColor(np.asarray(result), cv2.COLOR_RGB2BGR)
        image_message = bridge.cv2_to_imgmsg(img, encoding="bgr8")
        template.append(image_message)
    #imgProcessingList=[]
    frameList=[]
    for frame,target in zip(imgProcessingList,template):
        newframe=Image()
        newframe.header=frame.header
        newframe.data=target.data
        newframe.encoding="bgr8"
        newframe.height=480
        newframe.width=640
        newframe.step=1920
        frameList.append(newframe)
    imgProcessingList=frameList
    processingImage.release()

def imuProcessByDTW():
    global imuList,imuCheckingList,data_train
    imu_buf.acquire()
    imuCheckingList=imuList
    imuList=[]
    imu_buf.release()
    processing.acquire()
    #scores=module.DTW.estimate_dtw(imuCheckingList,data_train)
    #if scores[0]>=2.0:
    #    imuCheckingList=data_replacing(imuCheckingList)
    processing.release()
    #print("DTW done! score : "+str(scores[0]))
    #print("DTW done!")

def data_replacing(template):
    global i,data_train
    print("当前片段为："+str(i))
    print("检测到IMU数据异常，开始替换数据")
    for item,target in zip(template,data_train):
        item.linear_acceleration.x=target[0]
        item.linear_acceleration.y=target[1]
        item.linear_acceleration.z=target[2]
    print("替换完成")
    return template


def publishImuByGroup():
    global imuSendList, imuCheckingList
    imuSendList=imuCheckingList
    imuCheckingList=[]
    imuSendList.reverse()
    while len(imuSendList)>0:
        imupub.publish(imuSendList.pop())
        imurate.sleep()

def publishImgByGroup():
    global imgProcessingList,imageSendList
    imageSendList=imgProcessingList
    imgProcessingList=[]
    imageSendList.reverse()
    while len(imageSendList)>0:
        imgpub.publish(imageSendList.pop())
        imgrate.sleep()

#
# def checkPublishStatus():
#     global i
#     FirstPublish=True
#     while True:
#         if FirstPublish:
#             if len(imuList)>=600 and len(imageList)>=60:
#                 print("first public")
#                 imgProcessTask = threading.Thread(target=imgProcess)
#                 imgProcessTask.start()
#                 imgProcessTask.join()
#                 imuPubTask=threading.Thread(target=initialPublishImuByGroup)
#                 imgPubTask=threading.Thread(target=initialPublishImgByGroup)
#                 imuPubTask.start()
#                 imgPubTask.start()
#                 imuPubTask.join()
#                 imgPubTask.join()
#                 FirstPublish=False
#                 i=i+1
#                 print("first public done")
#         else:
#             if  len(imuList)>=300  and len(imageList)>=30:
#                 imuCheckTask=threading.Thread(target=imuProcessByDTW)
#                 imgProcessTask=threading.Thread(target=imgProcess)
#
#                 imuCheckTask.start()
#                 imgProcessTask.start()
#
#                 imuCheckTask.join()
#                 imgProcessTask.join()
#                 print("imu & image processed")
#
#                 imuPubTask=threading.Thread(target=publishImuByGroup)
#                 imgPubTask=threading.Thread(target=publishImgByGroup)
#
#                 imuPubTask.start()
#                 imgPubTask.start()
#
#                 imuPubTask.join()
#                 imgPubTask.join()
#                 i=i+1
#                 print("imu & img publish done")

def checkPublishStatus():
    global i
    while True:
        if len(imuList) >= 300 and len(imageList) >= 30:
            imuCheckTask = threading.Thread(target=imuProcessByDTW)
            imgProcessTask = threading.Thread(target=imgCompressedProcess)

            imuCheckTask.start()
            imgProcessTask.start()

            imuCheckTask.join()
            imgProcessTask.join()
            print("imu & image processed")
            imuPubTask = threading.Thread(target=publishImuByGroup)
            imgPubTask = threading.Thread(target=publishImgByGroup)

            imuPubTask.start()
            imgPubTask.start()

            imuPubTask.join()
            imgPubTask.join()
            i = i + 1
            print("imu & img publish done")

def PubNoDelay(data):
    img_buf.acquire()
    cv_image = bridge.imgmsg_to_cv2(data, desired_encoding='bgr8')
    from PIL import Image
    frame = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
    result = module.lowlight_enhancement.lowlight(frame)
    img = cv2.cvtColor(np.asarray(result), cv2.COLOR_RGB2BGR)
    image_message = bridge.cv2_to_imgmsg(img, encoding="bgr8")
    image_message.header=data.header
    imgpub.publish(image_message)
    img_buf.release()



if __name__ == "__main__":

    rospy.init_node('FeatureProject',anonymous=True)
    initParameter()


    imusub = rospy.Subscriber(imu_topic, Imu,handleIMU, queue_size=20)
    imgsub=rospy.Subscriber(image_topic,Image,handleCompressedImage,queue_size=20)
    imupub=rospy.Publisher(imu_pub,Imu,queue_size=20)
    imgpub=rospy.Publisher(image_pub,Image,queue_size=20)

    imurate = rospy.Rate(300)
    imgrate = rospy.Rate(30)
    #
    # files = get_imu_record("imu_template")
    # for file in files:
    #     data_train = np.load(file)

    featureProject=threading.Thread(target=checkPublishStatus)
    featureProject.start()
    rospy.spin()

